"""
@Project    : cosmo-face
@Module     : resnet_face.py
@Author     : HuangJiWen[huangjiwen@haier.com]
@Created    : 2020/8/17 15:37
@Desc       : resnet18 centernet
"""

import torch.nn as nn

from models.center_face.resnet import BasicBlock, Bottleneck, ResNet18, ResNet50
from models.center_face.head import KeypointHead


class CenterResNetFace(nn.Module):

    def __init__(self, cfg):
        super(CenterResNetFace, self).__init__()

        if cfg["name"] == "resnet18":
            self.backbone_model = ResNet18(BasicBlock, [2, 2, 2, 2])
        elif cfg["name"] == "resnet50":
            self.backbone_model = ResNet50(Bottleneck, [3, 4, 6, 3])

        self.head_model = KeypointHead(cfg['in_channels'], cfg["head_conv"], cfg["num_keypoints"])   # 256, 64

    def forward(self, x):
        """x -> torch.Size([2, 3, 544, 544])"""
        x = self.backbone_model(x)  # torch.Size([2, 256, 136, 136])
        return self.head_model(x)
